import os
import json
import time
import gradio as gr
import pandas as pd
import xml.etree.ElementTree as ET
from datetime import datetime
import dateparser
import dashscope


def normalize_time(time_str):
    """英语时间规范化"""
    if not time_str or time_str.lower() in ['n/a', 'unknown']:
        return 'N/A'

    try:
        settings = {'DATE_ORDER': 'MDY'}
        dt = dateparser.parse(time_str, settings=settings)
        return dt.isoformat() if dt else time_str
    except:
        return time_str


def extract_events_from_story(story):
    """使用Chat方式的事件抽取"""
    try:
        messages = [
            {
                'role': 'system',
                'content': """You are a professional event extraction assistant. Extract structured events in JSON format with:
- head: Subject (person/entity)
- relation: Action/relationship (past tense)
- tail: Object/target
- time: Normalized timestamp (or N/A)"""
            },
            {
                'role': 'user',
                'content': f"""Story: {story}

Please respond ONLY with a valid JSON array of events. Example:
[
  {{"head": "Cinderella", "relation": "was forced to work for", "tail": "stepmother", "time": "after father's death"}},
  {{"head": "Fairy Godmother", "relation": "created", "tail": "glass slippers", "time": "midnight"}}
]"""
            }
        ]

        response = dashscope.Generation.call(
            model='qwen-max-2025-01-25',
            api_key=os.getenv('DASHSCOPE_API_KEY'),
            messages=messages,
            result_format='message'
        )

        # 解析响应
        json_str = response.output.choices[0].message.content
        events = json.loads(json_str)

        # 规范化时间
        for event in events:
            event['time'] = normalize_time(event.get('time', 'N/A'))

        return events
    except Exception as e:
        print(f"API Error: {str(e)}")
        return []


def save_events_to_xml(events, filename):
    """保存事件到XML"""
    root = ET.Element('story_events')

    # 元数据
    meta = ET.SubElement(root, 'metadata')
    ET.SubElement(meta, 'created').text = datetime.now().isoformat()
    ET.SubElement(meta, 'language').text = 'en'
    ET.SubElement(meta, 'source').text = filename

    # 事件数据
    events_elem = ET.SubElement(root, 'events')
    for event in events:
        event_elem = ET.SubElement(events_elem, 'event')
        ET.SubElement(event_elem, 'head').text = event.get('head', '')
        ET.SubElement(event_elem, 'relation').text = event.get('relation', '')
        ET.SubElement(event_elem, 'tail').text = event.get('tail', '')
        ET.SubElement(event_elem, 'time').text = event.get('time', 'N/A')

    # 格式化XML
    def indent(elem, level=0):
        spacer = "  "
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = "\n" + spacer * (level + 1)
            for child in elem:
                indent(child, level + 1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = "\n" + spacer * level
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = "\n" + spacer * level

    indent(root)

    os.makedirs("stories", exist_ok=True)
    full_path = os.path.join("stories", filename)
    ET.ElementTree(root).write(full_path, encoding='utf-8', xml_declaration=True)
    return full_path


def process_story(story_text):
    """处理函数"""
    start_time = time.time()

    try:
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"story_{timestamp}.xml"

        events = extract_events_from_story(story_text)

        if not events:
            return pd.DataFrame(), "⚠️ No events extracted", ""

        saved_path = save_events_to_xml(events, filename)

        df = pd.DataFrame(events)
        process_time = f"⏱️ Process time: {time.time() - start_time:.2f}s"

        return df, process_time, saved_path

    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        return pd.DataFrame(), error_msg, ""


# Gradio界面
with gr.Blocks(theme=gr.themes.Soft(), title="Story Event Extractor") as app:
    gr.Markdown("# 📖 English Story Event Extractor")
    gr.Markdown("Analyze stories to extract structured events")

    with gr.Row():
        input_text = gr.Textbox(label="Input Story", lines=8,
                                placeholder="Paste your English story here...")
        btn = gr.Button("Analyze", variant="primary")

    output_table = gr.Dataframe(
        label="Extracted Events",
        headers=["Subject", "Relation", "Object", "Time"],
        datatype=["str", "str", "str", "str"],
        interactive=False
    )
    process_info = gr.Textbox(label="Processing Info")
    file_output = gr.Textbox(label="Saved Path")

    gr.Examples(
        examples=[[
            "Cinderella lived with her cruel stepmother after her father's death. "
            "On the night of the royal ball, the Fairy Godmother transformed a pumpkin into a carriage. "
            "Cinderella danced with the prince until midnight when the spell broke."
        ]],
        inputs=input_text,
        outputs=[output_table, process_info, file_output],
        fn=process_story,
        cache_examples=False
    )

    btn.click(
        fn=process_story,
        inputs=input_text,
        outputs=[output_table, process_info, file_output]
    )

if __name__ == "__main__":
    app.launch(server_name="127.0.0.1", server_port=7860)